import torch
import  torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from model import *


train_data=torchvision.datasets.CIFAR10("./dataset",train=True,transform=torchvision.transforms.ToTensor(),download=True)
test_data=torchvision.datasets.CIFAR10("./dataset",train=False,transform=torchvision.transforms.ToTensor(),download=True)

train_data_size = len(train_data)
test_data_size = len(test_data)
print(f"训练数据集的长度为{train_data_size}")
print(f"测试数据集的长度为{test_data_size}")

# 添加tensorboard
writer=SummaryWriter("./logs")

train_dataloader=DataLoader(train_data,batch_size=64)
test_dataloader=DataLoader(test_data,batch_size=64)

jjw=JJw()

# 损失函数
loss_fun = nn.CrossEntropyLoss()

# 优化器
learn_rate=0.01
optimizer=torch.optim.SGD(params=jjw.parameters(),lr=learn_rate)

# 设置训练网络的一些参数
total_train_step=0 # 记录训练的次数
total_test_step=0  # 记录测试的次数
total_accuracy=0 # 准确度
epoch=10 # 训练的轮数

for i in range(epoch):
    print(f"第{i+1}轮训练开始")

    # 训练步骤开始
    for data in train_dataloader:
        imgs,targets=data
        outputs=jjw(imgs)
        loss=loss_fun(outputs,targets)
        # 优化器优化模型 参数
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # 记录训练次数
        total_train_step+=1
        if total_train_step%100==0:
            print(f"训练次数={total_train_step},Loss={loss}")
            writer.add_scalar("train_loss",loss.item(),total_train_step)

    # 测试步骤开始
    total_test_loss=0
    with torch.no_grad():
        for data in test_dataloader:
            imgs,targets=data
            outputs=jjw(imgs)
            loss=loss_fun(outputs,targets)
            total_test_loss+=loss.item()
            accuracy=sum(outputs.argmax(1)==targets)
            total_accuracy+=accuracy

    print(f"整体测试集上的loss{total_test_loss}")
    print(f"整体测试集上的正确率{total_accuracy/test_data_size}")
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    writer.add_scalar("test_accuracy",total_test_loss, total_test_step,total_test_step)
    total_test_step += 1

    torch.save(jjw,f"jjw_{i}.pth")
    print("模型已保存")

writer.close()




